/*
 * normal_distribution.h
 *
 *  Created on: Feb 2, 2012
 *      Author: kkb110
 */

#ifndef NORMAL_DISTRIBUTION_H_
#define NORMAL_DISTRIBUTION_H_

#include "common_includes.h"

template<typename T, size_t N>
class NormalDistribution{
public:
	Eigen::Matrix<T, N, 1> mean;
	Eigen::Matrix<T, N, N> covariance;

	NormalDistribution(){
	}

	NormalDistribution(
			Eigen::Matrix<T, N, 1>& mean
			, Eigen::Matrix<T, N, N>& covariance){
		this->mean = mean;
		this->covariance = covariance;
	}

	T probability_density(const Eigen::Matrix<T, N, 1>& point) const{
		if (!covariance.matrix().ldlt().isPositive()){
			return 0;
		}
		Eigen::Matrix<T, N, N> covariance_inverse = covariance.inverse();
		Eigen::Matrix<T, N, 1> point_offset = point - mean;
		T tmp0 = 1
				/ (pow(2.0 * M_PI, N / 2.0)
						* sqrt(abs(covariance.matrix().determinant())));
		T tmp1 = point_offset.matrix().transpose() * covariance_inverse
				* point_offset.matrix(); //exp(-0.5);
		return tmp0 * exp(-0.5 * tmp1);
	}

	template<typename RandomNumberEngine>
	vector<Eigen::Matrix<T, N, 1>> generate_points(
			size_t n,
			RandomNumberEngine random_number_engine) const{
		Eigen::EigenSolver<decltype(covariance)> a(covariance);
		Eigen::Matrix<T, N, 1> b = a.eigenvalues().real().array().sqrt().sqrt();// TODO sqrt??
		Eigen::Matrix<T, N, N> c = b.matrix().asDiagonal()
				* a.eigenvectors().transpose().real();
		Eigen::Matrix<T, N, N> d = c.transpose() * c;
		normal_distribution<double> normal_distribution_generator(0, 1);

		Eigen::Matrix<T, Eigen::Dynamic, N> r(n, N);
		for (size_t i = 0; i < n; ++i){
			for (size_t j = 0; j < N; ++j){
				r(i, j) = normal_distribution_generator(random_number_engine);
			}
		}

		vector<Eigen::Matrix<T, N, 1>> result(n);
		Eigen::Matrix<T, N, Eigen::Dynamic> result_matrix =
				(r * d).transpose().colwise() + mean;
		for (size_t i = 0; i < n; ++i){
			result[i] = result_matrix.col(i);
		}
		return result;
	}

	void save(string file_name){
		save_eigen(mean, file_name + "_mean");
		save_eigen(covariance, file_name + "_covariance");
	}

	int progress(
			void *instance,
			const lbfgsfloatval_t *x,
			const lbfgsfloatval_t *g,
			const lbfgsfloatval_t fx,
			const lbfgsfloatval_t xnorm,
			const lbfgsfloatval_t gnorm,
			const lbfgsfloatval_t step,
			int n,
			int k,
			int ls){
		cout << k << " gnorm : " << gnorm << endl;

		return 0;
	}

	lbfgsfloatval_t evaluate(
			void* instance,
			const lbfgsfloatval_t* x,
			lbfgsfloatval_t* g,
			const int n,
			const lbfgsfloatval_t step){
		const size_t NUM_TOTAL = N + (N + 1) * N / 2;
		NormalDistribution<Number<T, NUM_TOTAL> ,N> model;
		for (size_t i = 0; i < N; ++i){
			model.mean(i) = Number<T, NUM_TOTAL>(i, x[i]);
		}
		for (size_t i = 0, index = N; i < N; ++i){

			for (size_t j = i; j < N; ++j){
				model.covariance(i, j) = Number<T, NUM_TOTAL>(index, x[index]);
				model.covariance(j, i) = model.covariance(i, j);
				++index;
			}
		}
		Number<T, NUM_TOTAL> result0 = 0;
		Number<T, NUM_TOTAL> result1 = 1;
		for (const Eigen::Matrix<T, N, 1>& shot : *fit_shots){
			Eigen::Matrix<Number<T, NUM_TOTAL> ,N, 1> convert;
			for (unsigned i = 0; i < N; ++i){
				convert(i) = (Number<T, NUM_TOTAL> ) shot(i);
			}
			result0 += log(model.probability_density(convert));
//			result1 *= model.probability_density(convert);
		}

		auto result = result0 + result1;
//		auto result = result0;
		result = -result;
		for (size_t i = 0; i < NUM_TOTAL; ++i){
			g[i] = result[i];
		}
		return (lbfgsfloatval_t) result.var;
	}

	static int progress_wrapper(
			void *instance,
			const lbfgsfloatval_t *x,
			const lbfgsfloatval_t *g,
			const lbfgsfloatval_t fx,
			const lbfgsfloatval_t xnorm,
			const lbfgsfloatval_t gnorm,
			const lbfgsfloatval_t step,
			int n,
			int k,
			int ls){
		return static_cast<NormalDistribution*>(instance)->progress(instance, x,
				g, fx, xnorm, gnorm, step, n, k, ls);
	}

	static lbfgsfloatval_t evaluate_wrapper(
			void* instance,
			const lbfgsfloatval_t* x,
			lbfgsfloatval_t* g,
			const int n,
			const lbfgsfloatval_t step){
		return static_cast<NormalDistribution*>(instance)->evaluate(instance, x,
				g, n, step);

	}

private:
	const vector<Eigen::Matrix<T, N, 1>>* fit_shots;
public:
	void fit(const vector<Eigen::Matrix<T, N, 1>>& shots){
		fit_shots = &shots;
		lbfgsfloatval_t fx;
		const size_t NUM_TOTAL = N + (N + 1) * N / 2;
		lbfgsfloatval_t * x = lbfgs_malloc(NUM_TOTAL);
		for (int i = 0; i < NUM_TOTAL; ++i){
			x[i] = 0;
		}
		for (int i = 0, index = N; i < N; ++i){
			x[index] = 1;
			index += N - i;
		}

		lbfgs_parameter_t param;
		lbfgs_parameter_init(&param);
//		param.epsilon *= param.epsilon;
//		param.min_step *= param.min_step;
//		param.m = 1;
		auto error_code = lbfgs(NUM_TOTAL, x, &fx, evaluate_wrapper,
				progress_wrapper, this, &param);
		cout << "error code : " << error_code << endl;

		for (size_t i = 0; i < N; ++i){
			mean(i) = x[i];
		}
		for (size_t i = 0, index = N; i < N; ++i){
			for (size_t j = 0; j < i; ++j){
				covariance(i, j) = covariance(j, i);
			}
			for (size_t j = i; j < N; ++j){
				covariance(i, j) = x[index];
				++index;
			}
		}
		lbfgs_free(x);
	}

};

#endif /* NORMAL_DISTRIBUTION_H_ */
